﻿#define STB_IMAGE_IMPLEMENTATION
#define STB_IMAGE_WRITE_IMPLEMENTATION
#define _CRT_SECURE_NO_DEPRECATE
#define num_threads 16
#define BLOCK_SIZE 20

#include "stb_image.h"
#include "stb_image_write.h"

#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <string.h>
#include <omp.h>

/**
 * @brief 读取图像文件（支持JPG、PNG、BMP等）
 * @param filename 图像文件路径
 * @param width 返回图像宽度的指针
 * @param height 返回图像高度的指针
 * @param channels 返回图像通道数的指针（如 RGB 为 3）
 * @return 返回图像数据的 unsigned char* 指针（需要使用 stbi_image_free() 释放），失败时返回 NULL
 */
unsigned char* load_image(const char* filename, int* width, int* height, int* channels) {
    unsigned char* data = stbi_load(filename, width, height, channels, 0);
    if (!data) {
        fprintf(stderr, "无法读取图像文件: %s\n", filename);
        return NULL;
    }
    return data;
}

/**
 * @brief 保存图像数据为 PNG 文件（灰度图或彩色图）
 * @param filename 要保存的文件名（应以 .png, .jpg, .bmp 结尾）
 * @param data 图像数据指针（灰度图为每像素1字节）
 * @param width 图像宽度
 * @param height 图像高度
 * @param channels 图像通道数（灰度=1，RGB=3）
 * @return 成功返回1，失败返回0
 */
int save_image(const char* filename, const unsigned char* data,
    int width, int height, int channels) {
    int success = 0;

    if (strstr(filename, ".png")) {
        success = stbi_write_png(filename, width, height, channels, data, width * channels);
    }
    else if (strstr(filename, ".jpg") || strstr(filename, ".jpeg")) {
        success = stbi_write_jpg(filename, width, height, channels, data, 95); // 95为质量
    }
    else if (strstr(filename, ".bmp")) {
        success = stbi_write_bmp(filename, width, height, channels, data);
    }
    else {
        fprintf(stderr, "不支持的文件扩展名: %s\n", filename);
    }

    if (!success) {
        fprintf(stderr, "图像保存失败: %s\n", filename);
    }

    return success;
}


/**
 * @brief 对图像进行灰度变换（线性或对数）
 * @param data 原始图像数据（RGB 或 RGBA）
 * @param width 图像宽度
 * @param height 图像高度
 * @param channels 图像通道数（3 或 4）
 * @param mode 变换模式：0=线性，1=对数
 * @param a 参数 a
 * @param b 参数 b
 * @return 返回新的灰度图像数据（单通道），失败时返回 NULL。使用完毕后需 free()
 */
unsigned char* grayscale_transform(
    const unsigned char* data, int width, int height, int channels,
    int mode, float a, float b
) {
    if (channels < 3) return NULL;  // 至少需要RGB

    int size = width * height;
    unsigned char* gray = (unsigned char*)malloc(size);
    if (!gray) return NULL;
    int i;
    #pragma omp parallel for
    for (i = 0; i < size; ++i) {
        int index = i * channels;
        float r = data[index];
        float g = data[index + 1];
        float b0 = data[index + 2];

        // 灰度化公式（luminosity method）
        float gray_val = 0.299f * r + 0.587f * g + 0.114f * b0;

        // 应用灰度变换
        if (mode == 0) {
            // 线性变换
            gray_val = a * gray_val + b;
        }
        else {
            // 对数变换，避免 log(0)
            gray_val = a * logf(1.0f + gray_val) + b;
        }

        // 裁剪到 [0, 255]
        if (gray_val < 0) gray_val = 0;
        if (gray_val > 255) gray_val = 255;

        gray[i] = (unsigned char)(gray_val);
    }

    return gray;
}

// //分块尝试
//unsigned char* histogram_equalization_block(const unsigned char* gray, int width, int height) {
//    int size = width * height;
//    unsigned char* out = (unsigned char*)malloc(size);
//    if (!out) return NULL;
//
//    // 遍历每个块
//    for (int block_y = 0; block_y < height; block_y += BLOCK_SIZE) {
//        for (int block_x = 0; block_x < width; block_x += BLOCK_SIZE) {
//            // 计算当前块的实际大小（处理边界情况）
//            int block_width = (block_x + BLOCK_SIZE) < width ? BLOCK_SIZE : (width - block_x);
//            int block_height = (block_y + BLOCK_SIZE) < height ? BLOCK_SIZE : (height - block_y);
//            int block_size = block_width * block_height;
//
//            // 为当前块初始化直方图
//            int hist[256] = { 0 };
//
//            // 并行计算当前块的直方图
//            #pragma omp parallel for
//            for (int i = 0; i < block_size; ++i) {
//                int x = block_x + (i % block_width);
//                int y = block_y + (i / block_width);
//                unsigned char val = gray[y * width + x];
//             #pragma omp atomic
//                hist[val]++;
//            }
//
//            // 计算累积分布函数 (CDF)
//            int cdf[256] = { 0 };
//            cdf[0] = hist[0];
//            for (int i = 1; i < 256; ++i) {
//                cdf[i] = cdf[i - 1] + hist[i];
//            }
//
//            // 查找非零最小CDF值
//            int cdf_min = 0;
//            for (int i = 0; i < 256; ++i) {
//                if (cdf[i] != 0) {
//                    cdf_min = cdf[i];
//                    break;
//                }
//            }
//
//            // 构建查找表 (LUT)
//            unsigned char lut[256];
//            for (int i = 0; i < 256; ++i) {
//                lut[i] = (unsigned char)(255.0f * (cdf[i] - cdf_min) / (block_size - cdf_min) + 0.5f);
//            }
//
//            // 应用LUT到当前块
//            #pragma omp parallel for
//            for (int i = 0; i < block_size; ++i) {
//                int x = block_x + (i % block_width);
//                int y = block_y + (i / block_width);
//                out[y * width + x] = lut[gray[y * width + x]];
//            }
//        }
//    }
//
//    return out;
//}

/**
 * @brief 对灰度图像进行直方图均衡化
 * @param gray 原始灰度图像数据（每像素1字节）
 * @param width 图像宽度
 * @param height 图像高度
 * @return 返回新图像数据指针（需 free），失败返回 NULL
 */
unsigned char* histogram_equalization(const unsigned char* gray, int width, int height) {
    int size = width * height;
    int hist[256] = { 0 };

    // 用局部直方图合并避免写冲突
    int local_hists[num_threads][256];
    memset(local_hists, 0, sizeof(local_hists));

    int i;
    #pragma omp parallel
    {
        int tid = omp_get_thread_num();
        #pragma omp for
        for (i = 0; i < size; ++i) {
            unsigned char val = gray[i];
            local_hists[tid][val]++;
        }
    }

    // 合并局部直方图
    for (int t = 0; t < num_threads; ++t) {
        for (int i = 0; i < 256; ++i) {
            hist[i] += local_hists[t][i];
        }
    }

    // 累积分布函数
    int cdf[256] = { 0 };
    cdf[0] = hist[0];
    for (int i = 1; i < 256; ++i) {
        cdf[i] = cdf[i - 1] + hist[i];
    }

    // 查找非零最小CDF值
    int cdf_min = 0;
    for (int i = 0; i < 256; ++i) {
        if (cdf[i] != 0) {
            cdf_min = cdf[i];
            break;
        }
    }

    // LUT 构建
    unsigned char lut[256];
    for (int i = 0; i < 256; ++i) {
        lut[i] = (unsigned char)(255.0f * (cdf[i] - cdf_min) / (size - cdf_min) + 0.5f);
    }

    // 应用 LUT（并行化）
    unsigned char* out = (unsigned char*)malloc(size);
    if (!out) return NULL;

    int j;
    #pragma omp parallel for
    for (j = 0; j < size; ++j) {
        out[j] = lut[gray[j]];
    }

    return out;
}

/**
 * @brief 图像均值滤波（支持灰度或 RGB 图像）
 * @param input 输入图像数据（每像素有 channels 个字节）
 * @param width 图像宽
 * @param height 图像高
 * @param channels 通道数（1=灰度, 3=RGB）
 * @param kernel_size 核大小（必须为奇数）
 * @return 平滑后的新图像（需 free），失败返回 NULL
 */
unsigned char* mean_filter(const unsigned char* input, int width, int height, int channels, int kernel_size) {
    
    if (kernel_size % 2 == 0 || kernel_size < 1) {
        fprintf(stderr, "kernel_size 必须为奇数且 >= 1\n");
        return NULL;
    }

    int half = kernel_size / 2;
    int size = width * height * channels;
    unsigned char* output = (unsigned char*)malloc(size);
    if (!output) return NULL;


    int y, x, c, ky, kx, sum ,count;
    #pragma omp parallel for private(x, c, ky, kx, sum, count)
    for (y = 0; y < height; ++y) {
        for (x = 0; x < width; ++x) {
            for (c = 0; c < channels; ++c) {
                sum = 0;
                count = 0;

                for (ky = -half; ky <= half; ++ky) {
                    for (kx = -half; kx <= half; ++kx) {
                        int px = x + kx;
                        int py = y + ky;
                        if (px >= 0 && px < width && py >= 0 && py < height) {
                            int idx = (py * width + px) * channels + c;
                            sum += input[idx];
                            count++;
                        }
                    }
                }
                output[(y * width + x) * channels + c] = (unsigned char)(sum / count);
            }
        }
    }

    return output;
}


int test() {
    omp_set_num_threads(num_threads);
    const char* filename = "example.jpg";
    int width, height, channels;

    unsigned char* img = load_image(filename, &width, &height, &channels);
    if (!img) return 1;

    // 灰度变换：线性模式，a=1, b=0
    unsigned char* gray = grayscale_transform(img, width, height, channels, 1,1.0f,0.0f);
    save_image("gray_output.bmp", gray, width, height, 1);

    unsigned char* equalized = histogram_equalization(gray, width, height);
    save_image("equalized.bmp", equalized, width, height, 1);


    unsigned char* out_rgb = mean_filter(img, width, height, 3, 5);
    save_image("rgb_smooth.bmp", out_rgb, width, height, 5);

    free(out_rgb);
    free(equalized);
    free(gray);
    stbi_image_free(img);
    return 0;
}

int main(int argc, char* argv[]) {
    omp_set_num_threads(num_threads);
   
    
    if (argc < 4) {
        printf("用法:\n");
        printf("  灰度增强:\n    c_image_tools_multi.exe -g -l/-e a b example.jpg output.bmp\n");
        printf("  直方图均衡:\n    c_image_tools_multi.exe -h example.jpg output.bmp\n");
        printf("  直方图分块均衡:\n    c_image_tools_multi.exe -f example.jpg output.bmp\n");
        printf("  平滑:\n    c_image_tools_multi.exe -s example.jpg kernel_size output.bmp\n");
        return 1;
    }


    if (strcmp(argv[1], "-g") == 0) {
        if (argc != 7) {
            fprintf(stderr, "灰度增强参数错误\n");
            return 1;
        }

        int mode = (strcmp(argv[2], "-l") == 0) ? 0 : 1;
        float a = atof(argv[3]);
        float b = atof(argv[4]);
        const char* input_file = argv[5];
        const char* output_file = argv[6];

        int width, height, channels;
        unsigned char* img = load_image(input_file, &width, &height, &channels);
        if (!img) return 1;

        unsigned char* gray = grayscale_transform(img, width, height, channels, mode, a, b);
        if (gray) {
            save_image(output_file, gray, width, height, 1);
            free(gray);
        }
        stbi_image_free(img);
    }
    else if (strcmp(argv[1], "-h") == 0) {
        if (argc != 4) {
            fprintf(stderr, "直方图均衡参数错误\n");
            return 1;
        }

        const char* input_file = argv[2];
        const char* output_file = argv[3];

        int width, height, channels;
        unsigned char* img = load_image(input_file, &width, &height, &channels);
        if (!img) return 1;

        // 若是彩色图，先转灰度
        unsigned char* gray = (channels == 1) ? img : grayscale_transform(img, width, height, channels, 0, 1.0f, 0.0f);
        unsigned char* result = histogram_equalization(gray, width, height);
        if (result) {
            save_image(output_file, result, width, height, 1);
            free(result);
        }

        if (channels != 1) free(gray);
        stbi_image_free(img);
    }
    //else if (strcmp(argv[1], "-f") == 0) {
    //    if (argc != 4) {
    //        fprintf(stderr, "局部直方图均衡参数错误\n");
    //        return 1;
    //    }

    //    const char* input_file = argv[2];
    //    const char* output_file = argv[3];

    //    int width, height, channels;
    //    unsigned char* img = load_image(input_file, &width, &height, &channels);
    //    if (!img) return 1;

    //    // 若是彩色图，先转灰度
    //    unsigned char* gray = (channels == 1) ? img : grayscale_transform(img, width, height, channels, 0, 1.0f, 0.0f);
    //    unsigned char* result = histogram_equalization(gray, width, height);
    //    if (result) {
    //        save_image(output_file, result, width, height, 1);
    //        free(result);
    //    }

    //    if (channels != 1) free(gray);
    //    stbi_image_free(img);
    //}
   
    else if (strcmp(argv[1], "-s") == 0) {

        const char* input_file = argv[2];
        int kernel_size = atoi(argv[3]);
        const char* output_file = argv[4];

        if (kernel_size < 1 || kernel_size % 2 == 0) {
            fprintf(stderr, "kernel_size 必须是奇数且 >= 1\n");
            return 1;
        }

        int width, height, channels;
        unsigned char* img = load_image(input_file, &width, &height, &channels);
        if (!img) return 1;

        unsigned char* result = mean_filter(img, width, height, channels, kernel_size);
        if (result) {
            save_image(output_file, result, width, height, channels);
            free(result);
        }

        stbi_image_free(img);
    }
    else {
        fprintf(stderr, "未知命令选项: %s\n", argv[1]);
        return 1;
    }

    return 0;
}
